import pandas as pd
import numpy as np

from Indicators.macd import calculate_macd, macd_start_point
from strategies.macdStrategy import macdstartegy


# 假设data是包含股票数据的DataFrame，包含'high', 'low', 'close'列
# 计算SKDJ指标
def calculate_skdj(data, n=9, m1=3, m2=3):
    # 计算未成熟随机值（RSV）
    high = data['high']
    low = data['low']
    close = data['close']

    rsv = (close - low.rolling(n).min()) / (high.rolling(n).max() - low.rolling(n).min()) * 100

    # 计算K值和D值
    k = rsv.ewm(com=m1 - 1).mean()
    d = k.ewm(com=m2 - 1).mean()

    # 将K值和D值添加到DataFrame
    data['SKDJ_K'] = k
    data['SKDJ_D'] = d

    return data


# 应用SKDJ策略
def skdj_strategy(data, overbought_threshold=80, oversold_threshold=20):
    data = calculate_skdj(data)

    # 初始化信号列
    data['Signal'] = 0

    # 当SKDJ_D低于超卖阈值时，发出买入信号
    data.loc[data['SKDJ_D'] < oversold_threshold, 'Signal'] = 1

    # 当SKDJ_D高于超买阈值时，发出卖出信号
    data.loc[data['SKDJ_D'] > overbought_threshold, 'Signal'] = -1

    return data

if __name__ == '__main__':
    df = pd.read_csv('../data_info/bit_coin.csv')
    # 应用MACD和启动点检测
    df = calculate_skdj(df)
    df = skdj_strategy(df)
    df = calculate_macd(df)
    df = macd_start_point(df)
    print(df.head())
